#include <iostream>
#include <torch/extension.h>

// C++ API of my CUDA test programming
torch::Tensor my_cuda_test(torch::Tensor);

// C++ high level interface declarations & macro definition
#define CHECK_CUDA(x) TORCH_CHECK(x.device().is_cuda(), #x " must be a CUDA tensor")
#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x "must be contiguous")
#define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x)

torch::Tensor my_test(torch::Tensor in_tensor_1)
{
    // auto res_tensor = my_cuda_test(in_tensor_1);
    // return res_tensor;
    CHECK_INPUT(in_tensor_1);
    // my_cuda_test(in_tensor_1);
    // std::cout << "zgsb!!!" << std::endl;
    auto ret_val = my_cuda_test(in_tensor_1);
    return ret_val;
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m){
    m.def("lzs_test", &my_test, "function of my CUDA exercise");
}

